#include "UMManage.h"
#include "Lock.h"
#include "ring-buffer.h"

// 下面的静态变量 Enclave和非安全世界是不共享的

uint64_t um_addr = UINT64_MAX; // 起始虚拟地址
uint32_t um_size = UINT32_MAX;

const static uint32_t smcb_addr = 0; // smcb的起始地址相对偏移量
const static uint32_t smcb_size = 4096;

const static uint32_t e_mem_region_addr = smcb_addr + smcb_size; // E-BUFFER 的起始地址偏移量
uint32_t e_mem_region_size = UINT32_MAX;                         // E-BUFFER 的大小

uint32_t o_mem_region_addr = UINT32_MAX; // O-BUFFER 的起始地址相对偏移量
uint32_t o_mem_region_size = UINT32_MAX; // O-BUFFER 的大小

// smcb 结构体中的内容是共享的且占用的空间始终是 4096 字节，即第一个页面
struct smcb_t
{
    um_lock_t lock;

    volatile int64_t op_id;
    CircularBuffer e_buffer; // 只允许非安全世界可写
    CircularBuffer o_buffer; // 只允许安全世界可写
};

inline static void *vaddr(uint64_t addr)
{
    return (void *)((uint64_t)addr + um_addr);
}

inline static uint32_t offset(void *addr)
{
    return (uint32_t)((uint64_t)addr - um_addr);
}

inline static struct smcb_t *getSMCB()
{
    return (struct smcb_t *)vaddr(0);
}

inline static CircularBuffer *getEBuffer()
{
    return &getSMCB()->e_buffer;
}

inline static CircularBuffer *getOBuffer()
{
    return &getSMCB()->o_buffer;
}

#define BLOCK_SIZE_ALIGN_UP(size) (((size) + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE)
#define BLOCK_SIZE_ALIGN_DOWN(size) (((size) / BLOCK_SIZE) * BLOCK_SIZE)

inline static void __memcpy(void *dst, void *src, uint32_t size)
{
    char *dst_ = (char *)dst;
    char *src_ = (char *)src;
    for (uint32_t i = 0; i < size; i++)
    {
        dst_[i] = src_[i];
    }
    return;
}

inline static void __memset(void *addr, uint8_t val, uint32_t size)
{
    uint8_t *addr_ = (uint8_t *)addr;
    for (uint32_t i = 0; i < size; i++)
    {
        addr_[i] = val;
    }
}

void register_um(uint64_t addr_, uint32_t size_)
{
    um_addr = addr_;
    um_size = size_;

    e_mem_region_size = (um_size - e_mem_region_addr) / 2 / 4096 * 4096;

    o_mem_region_addr = um_size - e_mem_region_addr - e_mem_region_size;
    o_mem_region_size = um_size - o_mem_region_addr;
}

void init_um(bool is_obuffer)
{
    if (!is_obuffer)
    {
        initCircularBuffer(getEBuffer(),
                           e_mem_region_addr,
                           um_addr,
                           e_mem_region_size);
        TEST_PRINT("EBUFFER: buffer->data=%p,buffer->size=%u,buffer->front=%d,buffer->rear=%d,buffer->count=%d\n",
                   vaddr(getEBuffer()->data), getEBuffer()->buffer_size, getEBuffer()->front, getEBuffer()->rear, getEBuffer()->count);
    }
    else
    {
        initCircularBuffer(getOBuffer(),
                           o_mem_region_addr,
                           um_addr,
                           o_mem_region_size);
        TEST_PRINT("OBUFFER: buffer->data=%p,buffer->size=%u,buffer->front=%d,buffer->rear=%d,buffer->count=%d\n",
                   vaddr(getOBuffer()->data), getOBuffer()->buffer_size, getOBuffer()->front, getOBuffer()->rear, getOBuffer()->count);
    }

    TEST_PRINT("init_um: um_addr=0x%lx, um_size=%u, e_mem_region_addr=%p, e_mem_region_size=0x%x, o_mem_region_addr=%p, o_mem_region_size=0x%x\n",
               um_addr, um_size, vaddr(e_mem_region_addr), e_mem_region_size, vaddr(o_mem_region_addr), o_mem_region_size);
}

///////////////////////////////以下是call和ret/////////////////////////////////

static int64_t alloc_op_id()
{
    struct smcb_t *smcb = getSMCB();
    return __sync_fetch_and_add(&smcb->op_id, 1);
}

static int64_t async_call(uint32_t func_id,
                          void *param,
                          int param_len,
                          bool flag_ebuffer)
{
    CircularBuffer *buffer = flag_ebuffer ? getEBuffer() : getOBuffer();
    TEST_ASSERT_INFO(buffer != NULL, "buffer is NULL\n");
    TEST_ASSERT_INFO(buffer->buffer_size > 0, "buffer->buffer_size=%d\n", buffer->buffer_size);
    if (!isEnoughSpace(buffer, (uint32_t)sizeof(op_t) + (uint32_t)param_len))
    {
        return -1;
    }
    op_t op = {
        .op_id = alloc_op_id(),
        .func_id = func_id,
        .param_len = param_len,
        .type = UM_CALL,
    };

    TEST_ASSERT_INFO(op.op_id >= 0, "op.op_id=%ld\n", op.op_id);
    TEST_ASSERT_INFO(op.func_id >= 0, "op.func_id=%d\n", op.func_id);
    TEST_ASSERT_INFO(op.param_len >= 0, "op.param_len=%d\n", op.param_len);
    TEST_ASSERT_INFO(op.type == UM_CALL, "op.type=%d\n", op.type);

    if (0 == writeData(buffer, &op, param))
    {
        TEST_PRINT("flag_ebuffer=%d addr = %p, size = %ld, opid=%ld, type=%d\n", flag_ebuffer ? 1 : 0, vaddr(buffer->data) + buffer->rear, sizeof(op_t) + (size_t)param_len, op.op_id, UM_CALL);
        return -1;
    }

    return op.op_id;
}

static int async_ret(int64_t op_id,
                     void *ret,
                     int ret_len,
                     bool flag_ebuffer)
{
    CircularBuffer *buffer = flag_ebuffer ? getEBuffer() : getOBuffer();
    if (!isEnoughSpace(buffer, (uint32_t)sizeof(op_t) + (uint32_t)ret_len))
    {
        return -1;
    }

    op_t op = {
        .op_id = op_id,
        .func_id = 0,
        .param_len = ret_len,
        .type = UM_RET,
    };

    TEST_ASSERT_INFO(op.op_id >= 0, "op.op_id=%ld\n", op.op_id);
    TEST_ASSERT_INFO(op.func_id >= 0, "op.func_id=%d\n", op.func_id);
    TEST_ASSERT_INFO(op.param_len >= 0, "op.param_len=%d\n", op.param_len);
    TEST_ASSERT_INFO(op.type == UM_RET, "op.type=%d\n", op.type);

    if (0 == writeData(buffer, &op, ret))
    {
        TEST_PRINT("flag_ebuffer=%d addr = %p, size = %ld, opid=%ld, type=%d\n",
                   flag_ebuffer ? 1 : 0, vaddr(buffer->data) + buffer->rear,
                   sizeof(op_t) + (size_t)ret_len, op.op_id, UM_RET);
        return -1;
    }

    return 0;
}

int64_t async_ecall(uint32_t func_id,
                    void *param,
                    int param_len)
{
    // ecall 时，复制到 ebuffer 中
    return async_call(func_id, param, param_len, true);
}

int async_eret(int64_t op_id,
               void *ret,
               int ret_len)
{
    // eret 时，复制到 obuffer 中
    return async_ret(op_id, ret, ret_len, false);
}

int64_t async_ocall(uint32_t func_id,
                    void *param,
                    int param_len)
{
    // ocall 时，复制到 obuffer 中
    return async_call(func_id, param, param_len, false);
}

int async_oret(int64_t op_id,
               void *ret,
               int ret_len)
{
    // oret 时，复制到 ebuffer 中
    return async_ret(op_id, ret, ret_len, true);
}

static op_t find_call_or_ret(bool is_obuffer)
{
    op_t op = {
        .op_id = 0,
        .func_id = 0,
        .param_len = 0,
        .type = UM_UNUSED,
    };

    uint32_t count = 0;
    CircularBuffer *buffer = is_obuffer ? getOBuffer() : getEBuffer();
    // 原子读取 count
    count = __atomic_load_n(&buffer->count, __ATOMIC_SEQ_CST);

    TEST_ASSERT_INFO(count <= buffer->buffer_size, "count=%d, buffer->buffer_size=%d\n", count, buffer->buffer_size);
    TEST_ASSERT_INFO(count >= 0, "count=%d\n", count);

    if (count == 0)
    {
        return op;
    }

    // 读取数据头, 这里应该禁用用户直接访问ring-buffer，因为存在数据边界越过的问题
    op_t *_op = (op_t *)(vaddr(buffer->data) + buffer->front);
    if (buffer->front + sizeof(op_t) <= buffer->buffer_size)
    {
        __memcpy(&op, _op, sizeof(op_t));
    }
    else
    {
        // TEST_PRINT("op_t reachs the bound of buffer buffer->data=%p,buffer->size=%u,buffer->front=%d,buffer->rear=%d,buffer->count=%d\n",
        //            buffer->data, buffer->buffer_size, buffer->front, buffer->rear, atomic_read_uint32(&buffer->count));

        uint32_t spaceToEnd = buffer->buffer_size - buffer->front;
        __memcpy(&op, _op, spaceToEnd);
        __memcpy((void *)((uint64_t)&op + spaceToEnd), vaddr(buffer->data), sizeof(op_t) - spaceToEnd);
    }

    TEST_ASSERT_INFO(op.op_id >= 0, "op.op_id=%ld\n", op.op_id);
    TEST_ASSERT_INFO(op.func_id >= 0, "op.func_id=%d\n", op.func_id);
    TEST_ASSERT_INFO(op.param_len >= 0, "op.param_len=%d\n", op.param_len);
    TEST_ASSERT_INFO(op.type == UM_CALL || op.type == UM_RET,
                     "ebuffer=%d, buffer->data=%p,buffer->count=%u,buffer->front=%d,buffer->rear=%d,op.op_id=%ld,func_id=%d,param_len=%d,op.type=%d\n",
                     is_obuffer ? 0 : 1,
                     vaddr(buffer->data), buffer->count, buffer->front, buffer->rear,
                     op.op_id, op.func_id, op.param_len,
                     op.type);

    return op;
}

/**
 * @brief 读取数据, 会同时读取_op和data
 * @param dst 保存读取出的数据指针
 * @param _op 必须是 call_op_t/ret_op_t 类型的指针，保存读取出的header
 * @return 如果读取成果，返回值为true，否则返回false
 */
static bool read_call_or_ret(uint8_t *dst, op_t *op, bool is_ebuffer)
{
    TEST_ASSERT_INFO(dst != NULL, "dst is NULL\n");
    TEST_ASSERT_INFO(op != NULL, "op is NULL\n");
    TEST_ASSERT_INFO(op->param_len > 0, "op->param_len=%d\n", op->param_len);

    CircularBuffer *buffer = is_ebuffer ? getEBuffer() : getOBuffer();
    return readData(buffer, op, dst);
}

bool read_ecall_or_oret(uint8_t *dst, op_t *op)
{
    return read_call_or_ret(dst, op, true);
}

bool read_ocall_or_eret(uint8_t *dst, op_t *op)
{
    return read_call_or_ret(dst, op, false);
}

op_t find_ecall_or_oret()
{
    return find_call_or_ret(false);
}

op_t find_ocall_or_eret()
{
    return find_call_or_ret(true);
}

void drop_ebuffer_data()
{
    drop((CircularBuffer *)getEBuffer());
}

void drop_obuffer_data()
{
    drop((CircularBuffer *)getOBuffer());
}


uint32_t get_num_of_ecall_and_oret()
{
    return atomic_read_uint32(&(getEBuffer()->num_call_and_ret));
}

uint32_t get_num_of_ocall_and_eret()
{
    return atomic_read_uint32(&(getOBuffer()->num_call_and_ret));
}